'''
@Company: TWL
@Author: xue jian
@Email: xuejian@kanzhun.com
@Date: 2020-04-01 15:20:11
'''
# -*- coding: UTF-8 -*-
from subprocess import *
import random, time, json, sys
# from kafka import KafkaProducer
# from kafka import KafkaConsumer
# from kafka import TopicPartition
# import sys
# sys.path.append('../util')
# from feature_handler import get_fea_index, fea_code


class BatchTrain():
    def __init__(self, model_name, batch_size, shuffle_size, watch_num, neg_rate, process_num, dates, f_path, run_shell, net_conf, labels, random_sleep=0):
        # model_name += ('_' + str(time.strftime("%Y%m%d_%H%M%S", time.localtime())))
        self.model_name = model_name
        print("model_name = ", model_name)
        # exit(0)
        self.batch_size = batch_size
        self.process_num = process_num
        self.shuffle_size = shuffle_size
        self.watch_num = watch_num
        self.neg_rate = neg_rate
        self.run_shell = run_shell
        self.dates = dates
        self.f_path = f_path
        self.random_sleep = random_sleep
        self.net_conf = net_conf
        self.labels = labels
        self.shell_p = self.run_shell + " -log /data1/xuejian/sync/flash_train/conf/log.properties" + " -nets /data1/xuejian/sync/flash_train/conf/" + self.net_conf + " -mn " + self.model_name + " -w " + str(self.watch_num) + " -s " + str(self.shuffle_size) + " -b " + str(self.batch_size) + " -ne " + str(self.neg_rate) + " -log_out /data1/xuejian/sync/flash_train/log -labels " + self.labels
        print("shell = ", self.shell_p)



    def train(self):
        for date in self.dates:
            print(date)
            proc_dict = {}
            num_list = [0] * self.process_num
            print("num_list = ", num_list)
            for i in range(self.process_num):
                f_date = self.f_path + date + '/' + str(i)
                # f_date = self.f_path + date + '/' + str(sys.argv[2])
                if i != 0:
                    time.sleep(random.randint(0, self.random_sleep))
                proc = Popen('cat ' + f_date + ' | ' + self.shell_p, shell=True)
                # print('process start')
                proc_dict[i] = proc
            for index, p in proc_dict.items():
                p.wait()
                print('finished!! : ', index)

    def train2(self):

        proc = Popen(self.shell_p, shell=True, stdin=PIPE)
        print('process start')
        for date in self.dates:
            print(date)
            with open(self.f_path + date, 'rb') as f_date:
                print(self.f_path + date)
                line = f_date.readline()
                while line != b'':
                    if line == b'':
                        break
                    proc.stdin.write(line)
                    proc.stdin.flush()
                    line = f_date.readline()

        proc.stdin.close()
        proc.wait()


if __name__ == '__main__':
    dates = []
    # dates.extend(["2019-09-" + str(i) for i in range(20, 31)])
    # dates.extend(["2019-10-0" + str(i) for i in range(1, 10)])
    # dates.extend(["2019-10-" + str(i) for i in range(21, 22)])
    # dates.extend(["2019-10-0" + str(i) for i in range(8, 10)])
    # dates.extend(["2019-10-" + str(i) for i in range(10, 31)])
    # dates.extend(["2019-10-" + str(i) for i in range(21, 32)])
    # dates.extend(["2019-11-0" + str(i) for i in range(0, 10)])
    dates.extend(["2019-12-" + str(i) for i in range(16, 19)])

    batch_train = BatchTrain(sys.argv[1], 2000, 40000, 200000, 0.72, 6, dates, '/data2/training_data/recall_fid_cut/', '/data2/bin/flash/flash_train', 'base_recall.net', 'one', 14)
    batch_train.train()